package com.example.toy.core.reader;

import com.example.toy.core.dataset.BigCollection;
import com.example.toy.core.dataset.Pair;
import com.example.toy.core.filter.ToyFilter;
import com.example.toy.core.record.ToyRecord;
import com.example.toy.core.table.ToyTableMeta;
import org.apache.hadoop.fs.Path;

import java.util.List;
import java.util.stream.Collectors;

/**
 * DefaultToyPartitionReader
 *
 * @author shenbing
 * @date 2021/12/20
 */
public class DefaultToyPartitionReader implements ToyPartitionReader {
  private ToyTableMeta tableMeta;
  private String partitionKey;
  private Path partitionPath;

  public DefaultToyPartitionReader(ToyTableMeta tableMeta, String partitionKey) {
    this.tableMeta = tableMeta;
    this.partitionKey = partitionKey;
    if (!this.tableMeta.getPartitionPathMap().containsKey(partitionKey)) {
      throw new IllegalArgumentException("Partition not exists: " + partitionKey);
    }

    this.partitionPath = tableMeta.getPartitionPathMap().get(partitionKey);
  }

  @Override
  public String getPartitionKey() {
    return this.partitionKey;
  }

  @Override
  public BigCollection<ToyRecord> read(ToyFilter... filters) {
    ToyFileReader baseFileReader = baseFileReader();
    List<ToyFileReader> deltaFileReaders = deltaFileReaders();

    BigCollection<ToyRecord> baseRecords = baseFileReader.parRead(filters);
    List<BigCollection<ToyRecord>> deltaRecords =
        deltaFileReaders.stream().map(d -> d.parRead(filters)).collect(Collectors.toList());

    BigCollection<ToyRecord> records =
        deltaRecords.stream().reduce(baseRecords, (l, r) -> l.union(r));

    BigCollection<ToyRecord> merged =
        records.mapToPair(r -> new Pair<>(r.getKey(), r)).reduceByKey((l, r) -> l.merge(r));

    return merged;
  }

  public ToyFileReader baseFileReader() {
    List<Path> baseFile = tableMeta.getVersionFileView().getBaseFiles(this.partitionKey);
    if (baseFile.isEmpty()) {
      throw new IllegalStateException("Can not find base file of partition: " + partitionKey);
    }
    DefaultToyBaseFileReader reader = new DefaultToyBaseFileReader(this.tableMeta, baseFile.get(0));
    return reader;
  }

  public List<ToyFileReader> deltaFileReaders() {
    List<Path> deltaFiles = tableMeta.getVersionFileView().getDeltaFiles(this.partitionKey);
    List<ToyFileReader> deltaFileReaders =
        deltaFiles.stream()
            .map(d -> new DefaultToyDeltaFileReader(this.tableMeta, d))
            .collect(Collectors.toList());
    return deltaFileReaders;
  }
}
